import caffe
import numpy as np

class DatasetChoice(caffe.Layer):
    """
    Choice datset form coco or ai
    """

    def setup(self, bottom, top):
        """
        bottom 0,1,2,4 is coco_data,coco_label,ai_data,ai_label respectly
        """
        if len(bottom) != 4:
            raise Exception("Need two inputs to compute distance.")

        self.coco_flag = 0
        self.ai_flag = 1
        self.flags = [self.coco_flag,self.ai_flag]

    def reshape(self, bottom,top):
        assert bottom[0].count == bottom[2].count and bottom[1].count == bottom[3].count, 'coco dataset and ai datset shold have the same dimension of data and label, but there are is coco: {}, {}; ai: {},{}'.format(bottom[0].data.shape,bottom[1].data.shape,bottom[2].data.shape,bottom[3].data.shape)

        top[0].reshape(*bottom[0].data.shape)
        top[1].reshape(*bottom[1].data.shape)

    def forward(self,bottom,top):
        flag = np.random.choice(self.flags)        

        if  flag == self.coco_flag:
            #print 'from coco'
            top[0].data[...] = bottom[0].data
            top[1].data[...] = bottom[1].data
        else:
            #print 'from ai'
            top[0].data[...] = bottom[2].data
            top[1].data[...] = bottom[3].data

    def backward(self,top,propagate_down, bottom):
        """
        do not need backward
        """
        pass
            
